package com.gitee.wsl.flow

import com.gitee.wsl.ext.K
import com.gitee.wsl.ext.runBlocking
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope

import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.distinctUntilChanged
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.onCompletion
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.shareIn
import kotlinx.coroutines.plus
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import timber.log.Timber
import kotlin.coroutines.CoroutineContext

/**
 * A thread safe stateful flow that can be updated blocking and async with a lazy initial value provider.
 *
 * @param loggingTag will be prepended to logging tag, i.e. "$loggingTag:HD"
 * @param parentScope on which the update operations and callbacks will be executed on
 * @param coroutineContext used in combination with [CoroutineScope]
 * @param startValueProvider provides the first value, errors will be rethrown on [CoroutineScope]
 */
class DynamicStateFlow<T>(
    loggingTag: String? = null,
    parentScope: CoroutineScope,
    coroutineContext: CoroutineContext = parentScope.coroutineContext,
    private val onRelease: CoroutineScope.(T) -> Unit = {},
    private val startValueProvider: suspend CoroutineScope.() -> T,
) {
    private val lTag = loggingTag?.let { "$it:DSFlow" }?:""

    private val updateActions = MutableSharedFlow<Update<T>>(
        replay = Int.MAX_VALUE,
        extraBufferCapacity = Int.MAX_VALUE,
        onBufferOverflow = BufferOverflow.SUSPEND
    )
    private val valueGuard = Mutex()

    private val producer: Flow<State<T>> = channelFlow {
        var currentValue = valueGuard.withLock {
            Timber.tag(lTag).d("Providing startValue...")

            startValueProvider().also { startValue ->
                val initializer = Update<T>(onError = null, onModify = { startValue })
                send(State(value = startValue, updatedBy = initializer))
                Timber.tag(lTag).d("...startValue provided and emitted.")
            }
        }

        invokeOnClose {
            Timber.tag(lTag).d("invokeOnClose executing...")
            onRelease(currentValue)
            Timber.tag(lTag).d("internal channelFlow finished.")
        }

        updateActions.collect { update ->
            currentValue = valueGuard.withLock {
                try {
                    update.onModify(currentValue).also {
                        send(State(value = it, updatedBy = update))
                    }
                } catch (e: Exception) {
                    Timber.tag(lTag).d("Data modifying failed (onError=${update.onError}): $e")


                    if (update.onError != null) {
                        update.onError.invoke(e)
                    } else {
                        send(State(value = currentValue, error = e, updatedBy = update))
                    }

                    currentValue
                }
            }
        }
    }

    private val internalFlow = producer
        .onStart { Timber.tag(lTag).d( "Internal onStart") }
//        .onEach { value -> Timber.tag(lTag).d( "New value: $value") }
        .onCompletion { err ->
            when {
                err is CancellationException -> {
                    Timber.tag(lTag).d( "internal onCompletion() due to cancellation")
                }
                err != null -> {
                    Timber.tag(lTag).d( "internal onCompletion() due to error: ${err}")
                }
                else -> {
                    Timber.tag(lTag).d( "internal onCompletion()")
                }
            }
        }
        .shareIn(
            scope = parentScope + coroutineContext,
            replay = 1,
            started = SharingStarted.Lazily
        )

    val flow: Flow<T> = internalFlow
        .map { it.value }
        .distinctUntilChanged()

    suspend fun value() = flow.first()

    /**
     * Non blocking update method.
     * Gets executed on the scope and context this instance was initialized with.
     *
     * @param onError if you don't provide this, and exception in [onUpdate] will the scope passed to this class
     */
    fun updateAsync(
        onError: (suspend (Exception) -> Unit) = { throw it },
        onUpdate: suspend T.() -> T,
    ) {
        val update: Update<T> = Update(
            onModify = onUpdate,
            onError = onError
        )
        K.runBlocking { updateActions.emit(update) }
    }

    /**
     * Blocking update method
     * Gets executed on the scope and context this instance was initialized with.
     * Waiting will happen on the callers scope.
     *
     * Any errors that occurred during [action] will be rethrown by this method.
     */
    suspend fun updateBlocking(action: suspend T.() -> T): T {
        val update: Update<T> = Update(onModify = action)
        updateActions.emit(update)

        Timber.tag(lTag).d("Waiting for update.")
        val ourUpdate = internalFlow.first { it.updatedBy == update }
        Timber.tag(lTag).d( "Finished waiting, got $ourUpdate" )

        ourUpdate.error?.let { throw it }

        return ourUpdate.value
    }

    private data class Update<T>(
        val onModify: suspend T.() -> T,
        val onError: (suspend (Exception) -> Unit)? = null,
    )

    private data class State<T>(
        val value: T,
        val error: Exception? = null,
        val updatedBy: Update<T>,
    )
}